

#include "pmisAggregation.hpp"

#include <numeric>

#include "csrMatrix.hpp"
#include "mpi.h"
#include "mpiWrapper.hpp"
#include "unapMPI.hpp"
#include "utilityCommunicator.hpp"

#ifdef OMP
#include <omp.h>
#endif

namespace UNAP {

template <typename Cmpt>
PMISAggregation<Cmpt>::PMISAggregation(const CSRMatrix<Cmpt>& A)
    : eps_(0.08), A_(A), connection_(A.comm()), Ptent_(A.comm()) {
  label n = A.lRows();
  const label nProcs = A.comm()->getMySize();
  connection_.vtxdist() = A.vtxdist();
  if (nProcs > 1) {
    Array<label> state(n);
    Array<label> owner(n);
    buildConnStrength();
    nAggr_ = aggregate(state, owner);
    buildTentativeProlongation(nAggr_, state, owner);
  } else {
    RS_.build(A, Ptent_);
    nAggr_ = Ptent_.lCols();
  }
}

template <typename Cmpt>
void PMISAggregation<Cmpt>::buildConnStrength() {
  const label n = A_.lRows();
  const scalar* D = A_.diag().data();
  SPMVBuffers<Cmpt>& spmvbuf = A_.spMVBuffers();
  Array<scalar> Ddiag(spmvbuf.sbuf_.size());
  Array<scalar> Doffd(spmvbuf.rbuf_.size());

  for (ulabel i = 0; i < Ddiag.size(); i++)
    Ddiag[i] = D[spmvbuf.sind_[i] - A_.beginCol()];

  A_.exchange("S&R", Ddiag.data(), Doffd.data());
  A_.comm()->finishTask("S&R");

  connection_.lRows() = n;
  connection_.lCols() = n;
  connection_.beginRow() = A_.beginRow();
  connection_.beginCol() = A_.beginCol();
  connection_.val() = new char[A_.lnnz()];
  connection_.ptr() = new label[n + 1];
  char* val = connection_.val();
  label* ptr = connection_.ptr();
  for (label i = 0; i <= n; i++) ptr[i] = 0;

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; i++) {
    const scalar epsdiagi = eps_ * eps_ * D[i];

    for (label j = A_.ptr(i); j < A_.offDiagStart(i); j++) {
      const label c = A_.ind(j) - A_.beginCol();
      const Cmpt v = A_.val(j);
      if ((val[j] = (c == i || (epsdiagi * D[c] < v * v)))) ++ptr[i + 1];
    }
    unordered_map<label, std::pair<label, label> >& cmap = spmvbuf.rind_;
    for (label j = A_.offDiagStart(i); j < A_.ptr(i + 1); j++) {
      const label c = cmap[A_.ind(j)].second;
      const Cmpt v = A_.val(j);
      if ((val[j] = (epsdiagi * Doffd[c] < v * v))) ++ptr[i + 1];
    }
  }

  std::partial_sum(ptr, ptr + n + 1, ptr);
  connection_.lnnz() = ptr[n];
  connection_.ind() = new label[ptr[n]];
  label* col = connection_.ind();

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; i++) {
    label head = ptr[i];

    for (label j = A_.ptr(i); j < A_.ptr(i + 1); j++) {
      if (val[j]) col[head++] = A_.ind(j);
    }
  }
  connection_.splitOffDiag();
}

template <typename Cmpt>
CSRMatrix<char>* PMISAggregation<Cmpt>::matSquare(const CSRMatrix<char>& A) {
  const label Alrows = A.lRows();
  const label Alcols = A.lCols();
  const label Abegin = A.beginCol();
  const label Aend = Abegin + Alcols;
  const label *Aptr = A.ptr(), *Aind = A.ind();
  SPMVBuffers<char>& Aspmvbuf = A.spMVBuffers();

  const label* Aoffdiag = A.offDiagStart().data();
  CSRMatrix<char>* a_nbr = A.remoteRows(A);
  CSRMatrix<char>& A_nbr = *a_nbr;

  CSRMatrix<char>* C = new CSRMatrix<char>(A.comm());
  C->lRows() = Alrows, C->lCols() = Alcols;
  C->beginRow() = A.beginRow(), C->beginCol() = Abegin;
  C->ptr() = new label[Alrows + 1];
  C->vtxdist() = A.vtxdist();
  Array<label>& Coffdiag = C->offDiagStart();
  Coffdiag.resize(Alrows);
  label* const Cptr = C->ptr();
  Cptr[0] = 0;

  Array<label> remCols;
  label ind = A_nbr.lnnz();
  remCols.resize(ind + A.offDiagNnz());
  std::copy(A_nbr.ind(), A_nbr.ind() + ind, remCols.begin());
  for (label i = 0; i < Alrows; i++) {
    for (label j = Aoffdiag[i]; j < Aptr[i + 1]; j++) {
      remCols[ind++] = Aind[j];
    }
  }
  std::sort(remCols.begin(), remCols.end());
  remCols.erase(std::unique(remCols.begin(), remCols.end()), remCols.end());
  label count = 0;
  unordered_map<label, label> remIdx;
  for (ulabel i = 0; i < remCols.size(); i++) {
    const label c = remCols[i];
    if (c >= Abegin && c < Aend) continue;
    remIdx[c] = count++;
  }

#ifdef OMP
#pragma omp parallel
#endif
  {
    Array<label> locMarker(Alcols, -1);
    Array<label> remMarker(count, -1);
    label offDiagNnzCount = 0;
#ifdef OMP
#pragma omp for
#endif
    for (label ia = 0; ia < Alrows; ++ia) {
      label locCols = 0, remCols = 0;

      for (label ja = Aptr[ia]; ja < Aoffdiag[ia]; ja++) {
        const label ca = Aind[ja] - Abegin;
        for (label jb = Aptr[ca]; jb < Aptr[ca + 1]; jb++) {
          label cb = Aind[jb];
          if (cb >= Abegin && cb < Aend) {
            cb -= Abegin;
            if (locMarker[cb] != ia) {
              locMarker[cb] = ia;
              ++locCols;
            }
          } else {
            cb = remIdx[cb];
            if (remMarker[cb] != ia) {
              remMarker[cb] = ia;
              ++remCols;
            }
          }
        }
      }

      for (label ja = Aoffdiag[ia]; ja < Aptr[ia + 1]; ja++) {
        const label ca = Aspmvbuf.rind_[Aind[ja]].second;
        for (label jb = A_nbr.ptr(ca); jb < A_nbr.ptr(ca + 1); jb++) {
          label cb = A_nbr.ind(jb);
          if (cb >= Abegin && cb < Aend) {
            cb -= Abegin;
            if (locMarker[cb] != ia) {
              locMarker[cb] = ia;
              ++locCols;
            }
          } else {
            cb = remIdx[cb];
            if (remMarker[cb] != ia) {
              remMarker[cb] = ia;
              ++remCols;
            }
          }
        }
      }

      Coffdiag[ia] = locCols;
      Cptr[ia + 1] = locCols + remCols;
      offDiagNnzCount += remCols;
    }
    C->offDiagNnz() = offDiagNnzCount;
  }

  std::partial_sum(Cptr, Cptr + Alrows + 1, Cptr);
  for (label ia = 0; ia < Alrows; ia++) Coffdiag[ia] += Cptr[ia];
  C->lnnz() = Cptr[Alrows];
  C->ind() = new label[C->lnnz()];
  C->val() = new char[C->lnnz()];
  label* Cind = C->ind();

#ifdef OMP
#pragma omp parallel
#endif
  {
    Array<label> locMarker(Alcols, -1);
    Array<label> remMarker(count, -1);
#ifdef OMP
#pragma omp for
#endif
    for (label ia = 0; ia < Alrows; ++ia) {
      const label locBegin = Cptr[ia];
      const label remBegin = Coffdiag[ia];
      label locEnd = locBegin;
      label remEnd = remBegin;

      for (label ja = Aptr[ia]; ja < Aoffdiag[ia]; ja++) {
        const label ca = Aind[ja] - Abegin;
        for (label jb = Aptr[ca]; jb < Aptr[ca + 1]; jb++) {
          label cb = Aind[jb];
          if (cb >= Abegin && cb < Aend) {
            cb -= Abegin;
            if (locMarker[cb] < locBegin) {
              locMarker[cb] = locEnd;

              Cind[locEnd] = cb + Abegin;
              ++locEnd;
            }
          } else {
            const label gb = remIdx[cb];
            if (remMarker[gb] < remBegin) {
              remMarker[gb] = remEnd;

              Cind[remEnd] = cb;
              ++remEnd;
            }
          }
        }
      }

      for (label ja = Aoffdiag[ia]; ja < Aptr[ia + 1]; ja++) {
        const label ca = Aspmvbuf.rind_[Aind[ja]].second;
        for (label jb = A_nbr.ptr(ca); jb < A_nbr.ptr(ca + 1); jb++) {
          label cb = A_nbr.ind(jb);
          if (cb >= Abegin && cb < Aend) {
            cb -= Abegin;
            if (locMarker[cb] < locBegin) {
              locMarker[cb] = locEnd;

              Cind[locEnd] = cb + Abegin;
              ++locEnd;
            }
          } else {
            const label gb = remIdx[cb];
            if (remMarker[gb] < remBegin) {
              remMarker[gb] = remEnd;

              Cind[remEnd] = cb;
              ++remEnd;
            }
          }
        }
      }
    }
  }

  delete a_nbr;
  return C;
}

template <typename Cmpt>
label PMISAggregation<Cmpt>::aggregate(Array<label>& locState,
                                       Array<label>& locOwner) {
  enum { UNDONE = -2, DELETED = -1 };

  CSRMatrix<char>* S = matSquare(connection_);
  SPMVBuffers<char>& spmvbuf = S->spMVBuffers();
  const label* Sptr = S->ptr();
  const label* Sind = S->ind();
  const label* Soffdiag = S->offDiagStart().data();
  Communicator* comm = S->comm();
  const label* Aptr = connection_.ptr();
  const label* Aind = connection_.ind();
  const label* Aoffdiag = connection_.offDiagStart().data();
  const int myId = comm->getMyId();

  label nUndone = 0;
  const label ssize = spmvbuf.sbuf_.size();
  const label rsize = spmvbuf.rbuf_.size();

  Array<label> remState(rsize);
  Array<label> remOwner(rsize, -1);
  Array<label> sendState(ssize);
  Array<label> sendOwner(ssize);
  const label n = S->lRows();
  const label SbeginRow = S->beginRow();
  const label SbeginCol = S->beginCol();

  const label AbeginCol = connection_.beginCol();

  for (label i = 0; i < n; ++i) {
    const label w1 = Aoffdiag[i] - Aptr[i];
    const label w2 = Sptr[i + 1] - Soffdiag[i];
    if (w1 + w2 == 1) {
      locState[i] = DELETED;
      ++nUndone;
    } else {
      locState[i] = UNDONE;
    }
    locOwner[i] = -1;
  }
  nUndone = n - nUndone;

  for (label i = 0; i < ssize; ++i)
    sendState[i] = locState[spmvbuf.sind_[i] - SbeginRow];
  S->exchange("S&R", &sendState[0], &remState[0]);

  unordered_map<label, Array<label> > sendPts(spmvbuf.rranks_.size());
  Array<label> recvPts;
  Array<label> nbr;
  label naggr = 0;
  static const label tagCnt = 2001;
  static const label tagPts = 2002;
  Array<MPI_Request> sendCntReq(spmvbuf.rranks_.size());
  Array<MPI_Request> sendPtsReq(spmvbuf.rranks_.size());
  MPI_Comm mpiComm = comm->getMpiComm();
  comm->finishTask("S&R");

  while (true) {
    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i)
      sendPts[spmvbuf.rranks_[i]].clear();

    if (nUndone) {
      for (label i = 0; i < n; ++i) {
        if (locState[i] != UNDONE) continue;

        if (Sptr[i + 1] != Soffdiag[i]) {
          bool selectable = true;
          for (label j = Soffdiag[i]; j < Sptr[i + 1]; ++j) {
            const label js = Sind[j];
            const std::pair<label, label> c = spmvbuf.rind_[js];

            if (remState[c.second] == UNDONE && c.first > myId) {
              selectable = false;
              break;
            }
          }

          if (!selectable) continue;

          label id = naggr++;
          locOwner[i] = myId;
          locState[i] = id;
          --nUndone;

          for (label j = Aptr[i]; j < Aoffdiag[i]; ++j) {
            const label c = Aind[j] - AbeginCol;
            if (c != i) {
              if (locState[c] == UNDONE) --nUndone;
              locOwner[c] = myId;
              locState[c] = id;
            }
          }

          for (label j = Aoffdiag[i]; j < Aptr[i + 1]; ++j) {
            const label c = Aind[j];
            const std::pair<label, label> k = spmvbuf.rind_[c];

            remState[k.second] = id;

            sendPts[k.first].push_back(c);
            sendPts[k.first].push_back(id);
          }

          for (label j = Sptr[i]; j < Soffdiag[i]; ++j) {
            label c = Sind[j] - SbeginCol;
            if (c != i && locState[c] == UNDONE) {
              locOwner[c] = myId;
              locState[c] = id;
              --nUndone;
            }
          }

          for (label j = Soffdiag[i]; j < Sptr[i + 1]; ++j) {
            label c = Sind[j];
            const std::pair<label, label> k = spmvbuf.rind_[c];
            if (remState[k.second] == UNDONE) {
              remState[k.second] = id;
              sendPts[k.first].push_back(c);
              sendPts[k.first].push_back(id);
            }
          }

        } else {
          label id = naggr++;
          locOwner[i] = myId;
          locState[i] = id;
          --nUndone;

          nbr.clear();

          for (label j = Aptr[i]; j < Aoffdiag[i]; ++j) {
            label c = Aind[j] - AbeginCol;
            if (c != i && locState[c] != DELETED) {
              if (locState[c] == UNDONE) --nUndone;
              locOwner[c] = myId;
              locState[c] = id;
              nbr.push_back(c);
            }
          }

          for (ulabel m = 0; m < nbr.size(); m++) {
            const label k = nbr[m];
            for (label j = Aptr[k]; j < Aoffdiag[k]; ++j) {
              label c = Aind[j] - AbeginCol;
              if (c != k && locState[c] == UNDONE) {
                locOwner[c] = myId;
                locState[c] = id;
                --nUndone;
              }
            }
          }
        }
      }
    }

    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i) {
      label id = spmvbuf.rranks_[i];
      label npts = sendPts[id].size();
      MPI_Isend(&npts, 1 * sizeof(label), MPI_BYTE, id, tagCnt, mpiComm,
                &sendCntReq[i]);
      if (!npts) continue;
      MPI_Isend(&sendPts[id][0], npts * sizeof(label), MPI_BYTE, id, tagPts,
                mpiComm, &sendPtsReq[i]);
    }

    for (ulabel i = 0; i < spmvbuf.sranks_.size(); ++i) {
      const label id = spmvbuf.sranks_[i];
      label npts;
      MPI_Recv(&npts, 1 * sizeof(label), MPI_BYTE, id, tagCnt, mpiComm,
               MPI_STATUS_IGNORE);
      if (!npts) continue;
      recvPts.resize(npts);
      MPI_Recv(&recvPts[0], npts * sizeof(label), MPI_BYTE, id, tagPts, mpiComm,
               MPI_STATUS_IGNORE);

      for (label k = 0; k < npts; k += 2) {
        label c = recvPts[k] - SbeginCol;
        label r = recvPts[k + 1];
        if (locState[c] == UNDONE) --nUndone;
        locOwner[c] = id;
        locState[c] = r;
      }
    }

    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i) {
      const label id = spmvbuf.rranks_[i];
      label npts = sendPts[id].size();
      MPI_Wait(&sendCntReq[i], MPI_STATUS_IGNORE);
      if (!npts) continue;
      MPI_Wait(&sendPtsReq[i], MPI_STATUS_IGNORE);
    }

    for (label i = 0; i < ssize; ++i)
      sendState[i] = locState[spmvbuf.sind_[i] - SbeginRow];
    S->exchange("S&R", &sendState[0], &remState[0]);
    comm->finishTask("S&R");
    label count = nUndone;
    reduceSum(&count, comm);
    if (count == 0) break;
  }

  for (label i = 0; i < ssize; ++i)
    sendOwner[i] = locOwner[spmvbuf.sind_[i] - SbeginRow];
  S->exchange("S&R", &sendOwner[0], &remOwner[0]);
  Array<label> newId(naggr + 1, 0);
  comm->finishTask("S&R");
  for (label i = 0; i < n; ++i) {
    if (locOwner[i] == myId && locState[i] >= 0) newId[locState[i] + 1] = 1;
  }
  for (label i = 0; i < rsize; ++i) {
    if (remOwner[i] == myId && remState[i] >= 0) newId[remState[i] + 1] = 1;
  }
  std::partial_sum(newId.begin(), newId.end(), newId.begin());

  label count = naggr - newId.back();
  reduceSum(&count, comm);
  if (count > 0) {
    naggr = newId.back();

    for (label i = 0; i < n; ++i) {
      if (locOwner[i] == myId && locState[i] >= 0) {
        locState[i] = newId[locState[i]];
      }
    }
    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i)
      sendPts[spmvbuf.rranks_[i]].clear();

    for (unordered_map<label, std::pair<label, label> >::iterator iter =
             spmvbuf.rind_.begin();
         iter != spmvbuf.rind_.end(); ++iter) {
      const label c = iter->first;
      const std::pair<label, label> k = iter->second;

      if (remOwner[k.second] == myId && remState[k.second] >= 0) {
        sendPts[k.first].push_back(c);
        sendPts[k.first].push_back(newId[remState[k.second]]);
      }
    }

    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i) {
      label id = spmvbuf.rranks_[i];
      label npts = sendPts[id].size();
      MPI_Isend(&npts, 1 * sizeof(label), MPI_BYTE, id, tagCnt, mpiComm,
                &sendCntReq[i]);
      if (!npts) continue;
      MPI_Isend(&sendPts[id][0], npts * sizeof(label), MPI_BYTE, id, tagPts,
                mpiComm, &sendPtsReq[i]);
    }

    for (ulabel i = 0; i < spmvbuf.sranks_.size(); ++i) {
      label id = spmvbuf.sranks_[i];
      label npts;
      MPI_Recv(&npts, 1 * sizeof(label), MPI_BYTE, id, tagCnt, mpiComm,
               MPI_STATUS_IGNORE);

      if (!npts) continue;
      recvPts.resize(npts);
      MPI_Recv(&recvPts[0], npts * sizeof(label), MPI_BYTE, id, tagPts, mpiComm,
               MPI_STATUS_IGNORE);

      for (int k = 0; k < npts; k += 2) {
        label c = recvPts[k] - SbeginCol;
        label id = recvPts[k + 1];

        locState[c] = id;
      }
    }
    for (ulabel i = 0; i < spmvbuf.rranks_.size(); ++i) {
      const label id = spmvbuf.rranks_[i];
      label npts = sendPts[id].size();
      MPI_Wait(&sendCntReq[i], MPI_STATUS_IGNORE);
      if (!npts) continue;
      MPI_Wait(&sendPtsReq[i], MPI_STATUS_IGNORE);
    }
  }

  delete S;
  return naggr;
}

template <typename Cmpt>
void PMISAggregation<Cmpt>::buildTentativeProlongation(label naggr,
                                                       Array<label>& state,
                                                       Array<label>& owner) {
  enum { UNDONE = -2, DELETED = -1 };
  Communicator* comm = Ptent_.comm();
  const int nProcs = comm->getMySize();
  const int myId = comm->getMyId();
  const label n = connection_.lRows();
  label aggrSize[nProcs];
  comm->allGather("getGlobalDim", &naggr, sizeof(label), aggrSize,
                  sizeof(label));
  Array<label> aggrDist(nProcs + 1);
  label count = 0;
  comm->finishTask("getGlobalDim");
  for (label i = 0; i < nProcs; ++i) {
    aggrDist[i] = count;
    count += aggrSize[i];
  }
  aggrDist[nProcs] = count;

  Ptent_.lRows() = n;
  Ptent_.lCols() = naggr;
  Ptent_.beginRow() = connection_.beginRow();
  Ptent_.beginCol() = aggrDist[myId];
  Ptent_.vtxdist() = connection_.vtxdist();
  Ptent_.vtxdistCol() = aggrDist;

  Ptent_.ptr() = new label[n + 1];
  label* const Pptr = Ptent_.ptr();
  for (label i = 0; i <= n; ++i) Pptr[i] = 0;
  Array<label>& Poffdiag = Ptent_.offDiagStart();
  Poffdiag.resize(n);

  label offDiagNnzCount = 0;
#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; ++i) {
    if (state[i] == DELETED) continue;

    if (owner[i] == myId) {
      ++Poffdiag[i];
      ++Pptr[i + 1];
    } else {
      offDiagNnzCount++;
      ++Pptr[i + 1];
    }
  }

  std::partial_sum(Pptr, Pptr + n + 1, Pptr);
  for (label ia = 0; ia < n; ia++) Poffdiag[ia] += Pptr[ia];
  Ptent_.lnnz() = Pptr[n];
  Ptent_.ind() = new label[Pptr[n]];
  Ptent_.val() = new Cmpt[Pptr[n]];
  Ptent_.offDiagNnz() = offDiagNnzCount;
  label* const Pind = Ptent_.ind();
  Cmpt* const Pval = Ptent_.val();

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; ++i) {
    const label s = state[i];
    if (s == DELETED) continue;
    const label d = owner[i];
    Pind[Pptr[i]] = s + aggrDist[d];
    Pval[Pptr[i]] = static_cast<Cmpt>(1);
  }
}

template <typename Cmpt>
void PMISAggregation<Cmpt>::RugeStuben::build(const CSRMatrix<Cmpt>& A,
                                              CSRMatrix<Cmpt>& P) {
  const label n = A.lRows();
  static const scalar eps = 2 * std::numeric_limits<scalar>::epsilon();
  static const Cmpt zero = 0.;
  Array<char> cf(n, 'U');
  CSRMatrix<char> S(A.comm());

  connect(A, S, cf);
  splitcf(A, S, cf);

  label nc = 0;
  Array<label> cidx(n);
  for (label i = 0; i < n; ++i)
    if (cf[i] == 'C') cidx[i] = static_cast<label>(nc++);

  if (!nc) std::cerr << "Empty level" << std::endl;

  P.lRows() = n, P.lCols() = nc;
  P.rows() = n, P.cols() = nc;
  P.beginRow() = P.beginCol() = 0;
  P.setSymm(false);
  label* Pptr = new label[n + 1];
  P.ptr() = Pptr;
  for (label i = 0; i <= n; i++) Pptr[i] = 0;

  Array<Cmpt> Amin, Amax;

  if (ifTrunc_) {
    Amin.resize(n);
    Amax.resize(n);
  }

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; ++i) {
    if (cf[i] == 'C') {
      ++Pptr[i + 1];
      continue;
    }

    if (ifTrunc_) {
      Cmpt amin = zero, amax = zero;

      for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j) {
        if (!S.val(j) || cf[A.ind(j)] != 'C') continue;

        amin = std::min(amin, A.val(j));
        amax = std::max(amax, A.val(j));
      }

      Amin[i] = (amin *= epsTrunc_);
      Amax[i] = (amax *= epsTrunc_);

      for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j) {
        if (!S.val(j) || cf[A.ind(j)] != 'C') continue;

        if (A.val(j) < amin || amax < A.val(j)) ++Pptr[i + 1];
      }
    } else {
      for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j)
        if (S.val(j) && cf[A.ind(j)] == 'C') ++Pptr[i + 1];
    }
  }

  std::partial_sum(Pptr, Pptr + n + 1, Pptr);
  P.lnnz() = Pptr[n];
  label* Pind = new label[Pptr[n]];
  Cmpt* Pval = new Cmpt[Pptr[n]];
  P.ind() = Pind;
  P.val() = Pval;

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; ++i) {
    label row_head = Pptr[i];
    if (cf[i] == 'C') {
      Pind[row_head] = cidx[i];
      Pval[row_head] = (Cmpt)1.;
      continue;
    }

    Cmpt dia = zero;
    Cmpt a_num = zero, a_den = zero;
    Cmpt b_num = zero, b_den = zero;
    Cmpt d_neg = zero, d_pos = zero;

    for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j) {
      label c = A.ind(j);
      Cmpt v = A.val(j);
      if (c == i) {
        dia = v;
        continue;
      }

      if (v < zero) {
        a_num += v;
        if (S.val(j) && cf[c] == 'C') {
          a_den += v;
          if (ifTrunc_ && Amin[i] < v) d_neg += v;
        }
      } else {
        b_num += v;
        if (S.val(j) && cf[c] == 'C') {
          b_den += v;
          if (ifTrunc_ && v < Amax[i]) d_pos += v;
        }
      }
    }

    scalar cf_neg = 1;
    scalar cf_pos = 1;

    if (ifTrunc_) {
      if (std::abs(static_cast<Cmpt>(a_den - d_neg)) > eps)
        cf_neg = std::abs(a_den / (static_cast<Cmpt>(a_den - d_neg)));

      if (std::abs(static_cast<Cmpt>(b_den - d_pos)) > eps)
        cf_pos = std::abs(b_den / (static_cast<Cmpt>(b_den - d_pos)));
    }

    if (zero < b_num && std::abs(b_den) < eps) dia += b_num;

    scalar alpha =
        std::abs(a_den) > eps ? -cf_neg * std::abs(a_num / (dia * a_den)) : 0;
    scalar beta =
        std::abs(b_den) > eps ? -cf_pos * std::abs(b_num / (dia * b_den)) : 0;

    for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j) {
      label c = A.ind(j);
      Cmpt v = A.val(j);

      if (!S.val(j) || cf[c] != 'C') continue;
      if (ifTrunc_ && Amin[i] <= v && v <= Amax[i]) continue;

      Pind[row_head] = cidx[c];
      Pval[row_head] = (v < zero ? alpha : beta) * v;
      ++row_head;
    }
  }
  P.splitOffDiag();
}

template <typename Cmpt>
void PMISAggregation<Cmpt>::RugeStuben::connect(const CSRMatrix<Cmpt>& A,
                                                CSRMatrix<char>& S,
                                                Array<char>& cf) {
  const label n = A.lRows();
  const label nnz = A.lnnz();
  const scalar eps = 2 * std::numeric_limits<scalar>::epsilon();
  S.lRows() = S.lCols() = n;
  label* Sptr = new label[n + 1];
  label* Sind = new label[nnz];
  char* Sval = new char[nnz];
  S.ptr() = Sptr;
  S.ind() = Sind;
  S.val() = Sval;

  Sptr[0] = 0;

#ifdef OMP
#pragma omp parallel for
#endif
  for (label i = 0; i < n; ++i) {
    Sptr[i + 1] = 0;

    Cmpt a_min = 0.;

    for (label j = A.ptr(i); j < A.ptr(i + 1); j++)
      if (A.ind(j) != i) a_min = std::min(a_min, A.val(j));

    if (std::abs(a_min) < eps) {
      cf[i] = 'F';
      continue;
    }

    a_min *= epsStrong_;

    for (label j = A.ptr(i); j < A.ptr(i + 1); j++)
      Sval[j] = (A.ind(j) != i && A.val(j) < a_min);
  }

  for (label i = 0; i < nnz; ++i)
    if (Sval[i]) ++(Sptr[A.ind(i) + 1]);

  std::partial_sum(Sptr, Sptr + n + 1, Sptr);

  for (label i = 0; i < n; ++i)
    for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; ++j)
      if (Sval[j]) Sind[Sptr[A.ind(j)]++] = i;

  std::rotate(Sptr, Sptr + n, Sptr + n + 1);
  Sptr[0] = 0;
}

template <typename Cmpt>
void PMISAggregation<Cmpt>::RugeStuben::splitcf(const CSRMatrix<Cmpt>& A,
                                                const CSRMatrix<char>& S,
                                                Array<char>& cf) {
  const label n = A.lRows();
  Array<label> lambda(n);

  for (label i = 0; i < n; ++i) {
    label temp = 0;
    for (label j = S.ptr(i), e = S.ptr(i + 1); j < e; ++j)
      temp += (cf[S.ind(j)] == 'U' ? 1 : 2);
    lambda[i] = temp;
  }

  Array<label> ptr(n + 1, 0);
  Array<label> cnt(n, 0);
  Array<label> i2n(n);
  Array<label> n2i(n);
  for (label i = 0; i < n; ++i) ++ptr[lambda[i] + 1];

  std::partial_sum(ptr.begin(), ptr.end(), ptr.begin());

  for (label i = 0; i < n; ++i) {
    label lam = lambda[i];
    label idx = ptr[lam] + cnt[lam]++;
    i2n[idx] = i;
    n2i[i] = idx;
  }

  for (label top = n; top-- > 0;) {
    label i = i2n[top];
    label lam = lambda[i];

    if (lam == 0) {
      std::replace(cf.begin(), cf.end(), 'U', 'C');
      break;
    }

    --cnt[lam];

    if (cf[i] == 'F') continue;

    cf[i] = 'C';

    for (label j = S.ptr(i), e = S.ptr(i + 1); j < e; ++j) {
      label c = S.ind(j);

      if (cf[c] != 'U') continue;

      cf[c] = 'F';

      for (label aj = A.ptr(c), ae = A.ptr(c + 1); aj < ae; ++aj) {
        if (!S.val(aj)) continue;

        label ac = A.ind(aj);
        label lam_a = lambda[ac];

        if (cf[ac] != 'U' || static_cast<label>(lam_a) + 1 >= n) continue;

        label old_pos = n2i[ac];
        label new_pos = ptr[lam_a] + cnt[lam_a] - 1;

        n2i[i2n[old_pos]] = new_pos;
        n2i[i2n[new_pos]] = old_pos;

        std::swap(i2n[old_pos], i2n[new_pos]);

        --cnt[lam_a];
        ++cnt[lam_a + 1];
        ptr[lam_a + 1] = ptr[lam_a] + cnt[lam_a];

        lambda[ac] = lam_a + 1;
      }
    }

    for (label j = A.ptr(i), e = A.ptr(i + 1); j < e; j++) {
      if (!S.val(j)) continue;

      label c = A.ind(j);
      label lam = lambda[c];

      if (cf[c] != 'U' || lam == 0) continue;

      label old_pos = n2i[c];
      label new_pos = ptr[lam];

      n2i[i2n[old_pos]] = new_pos;
      n2i[i2n[new_pos]] = old_pos;

      std::swap(i2n[old_pos], i2n[new_pos]);

      --cnt[lam];
      ++cnt[lam - 1];
      ++ptr[lam];
      lambda[c] = lam - 1;
    }
  }
}

};  // namespace UNAP

template class UNAP::PMISAggregation<scalar>;
